import asyncio
import json
from multiprocessing.connection import Listener
import time
import os
from datetime import datetime
import uuid
from typing import Dict

from common import *
import tweb
from tweb import Request, Response, JSONBody

log = get_logger("server")
# 存储代理配置
agent_configs = {}
global_tokens = {}
g_agent_info = {}
# 存储转发规则和连接， key: listen_addr, value: {agent_id, target, server}
g_forward_rules = {}
# 持久化转发规则文件路径
RULES_FILE = os.path.join(os.path.dirname(__file__), 'rules.json')

class ServerSession:
    def __init__(self, listen, client_reader, client_writer):
        self.session_id = str(uuid.uuid4())[:8]  #固定8字节
        self.created_at = int(time.time())
        self.client_conn: str = ''
        self.target_conn: str = ''
        self.bytes_tx = 0
        self.bytes_rx = 0
        self.listen = listen
        self.type = 'local'
        self.client_reader = client_reader
        self.client_writer = client_writer
        # for local target
        self.target_reader = None
        self.target_writer = None
        # for remote target
        self.agent_writer = None
        self.client_closed = False
        self.target_connected = False
        self.tunnel_queue = asyncio.Queue()
        self.status = 'init'  # 'init', 'connecting', 'connected', 'closed'

    async def close_session(self):
        if self.status == 'closed':
            return
        self.status = 'closed'
        self.target_connected = False
        if self.client_writer and not self.client_closed:
            log.info(f'[{self.session_id}] closing client side')
            try:
                self.client_writer.close()
                await self.client_writer.wait_closed()
            except Exception as e:
                log.warning(f'[{self.session_id}] close client side error: {e}')
        if self.type == 'local':
            log.info(f'[{self.session_id}] closing local target side')
            try:
                if self.target_writer:
                    self.target_writer.close()
                    await self.target_writer.wait_closed()
            except Exception as e:
                log.warning(f'[{self.session_id}] close local target side error: {e}')
        elif self.agent_writer:
            log.info(f'[{self.session_id}] closing remote target side')
            try:
                self.tunnel_queue.put_nowait(b'') # notify loop to exit.
                self.agent_writer.write(encode_tlv(MSG_TYPE_SESSION_CLOSE, self.session_id.encode()))
                await self.agent_writer.drain()
            except Exception as e:
                log.warning(f'[{self.session_id}] close remote target side error: {e}')
        if self.session_id in g_sessions:
            del g_sessions[self.session_id]
            log.info(f"[{self.session_id}] session removed")

    async def local_forwarder(self,src_reader, src_writer, dst_reader, dst_writer, direction_lable):
        try:
            while True:
                # 从源读取数据
                data = await src_reader.read(4096)
                if not data:
                    log.info(f'[{self.session_id}] recv close msg. ({direction_lable} side)')
                    break
                # 写入目标
                dst_writer.write(data)
                await dst_writer.drain()
        except Exception as e:
            log.info(f"[{self.session_id}] forward error {direction_lable}: {e}")
        finally:
            await self.close_session()

    async def agent_to_client(self):
        """从隧道队列读取数据并转发到target"""
        try:
            while self.target_connected:
                data = await self.tunnel_queue.get()
                if not data:
                    log.info(f'[{self.session_id}][s->c] recv close msg. (tunnel side)')
                    break
                self.client_writer.write(data)
                # log.info(f'[{self.session_id}][s->c] transfer {len(data)} bytes')
                self.bytes_tx += len(data)
                await self.client_writer.drain()
                self.tunnel_queue.task_done()
        except Exception as e:
            log.info(f"[{self.session_id}][s->c] client write error: {e}")
        finally:
            log.info(f"[{self.session_id}][s->c] exit loop")
            await self.close_session()

    async def client_to_agent(self):
        """从client读取数据并转发到agent"""
        try:
            while self.target_connected:
                data = await self.client_reader.read(4096)
                if not data:
                    self.client_closed = True
                    log.info(f'[{self.session_id}][c->s] recv close msg. (client side)')
                    break
                # 携带session_id并封装为TLV格式发送到agent
                data = self.session_id.encode() + data
                tlv_data = encode_tlv(MSG_TYPE_SESSION_DATA, data)
                # log.info(f'[{self.session_id}][c->s] transfer {len(data)} bytes')
                self.bytes_rx += len(data)
                self.agent_writer.write(tlv_data)
                await self.agent_writer.drain()
        except Exception as e:
            log.info(f"[{self.session_id}][c->s] error: {e}")
        finally:
            log.info(f"[{self.session_id}][c->s] exit loop")
            await self.close_session()

    def start_data_forwarder(self):
        if self.type == 'local':
            asyncio.create_task(self.local_forwarder(self.client_reader, self.client_writer, self.target_reader, self.target_writer, 'client->target'))
            asyncio.create_task(self.local_forwarder(self.target_reader, self.target_writer, self.client_reader, self.client_writer, 'target->client'))
        else:
            asyncio.create_task(self.client_to_agent())
            asyncio.create_task(self.agent_to_client())

# 存储会话信息，key: session_id, value: ServerSession
g_sessions: Dict[str, ServerSession] = {}

def save_forward_rules():
    """将转发规则保存到JSON文件"""
    rules_data = {}
    for listen_addr, rule in g_forward_rules.items():
        # 只保存可序列化的字段
        rules_data[listen_addr] = {
            'listen': rule['listen'],
            'agent_id': rule['agent_id'],
            'target': rule['target']
        }
    try:
        with open(RULES_FILE, 'w') as f:
            json.dump(rules_data, f, indent=2)
        log.info(f"转发规则已保存到 {RULES_FILE}")
    except Exception as e:
        log.info(f"保存转发规则失败: {e}")

async def load_forward_rules():
    """从JSON文件加载转发规则并启动转发服务器"""
    log.info(f'loading rules. file={RULES_FILE}')
    if not os.path.exists(RULES_FILE):
        log.info("转发规则文件不存在，将使用空规则")
        return
    
    try:
        with open(RULES_FILE, 'r') as f:
            rules_data = json.load(f)
        
        for listen_addr, rule in rules_data.items():
            log.debug(f"load rule: {rule}")
            # 启动转发服务器
            success = await start_forward_server(
                listen_addr, 
                rule['agent_id'], 
                rule['target']
            )
            if not success:
                log.info(f"启动转发服务器失败: {rule}")
    except Exception as e:
        log.error(f"加载转发规则失败: {e}")
g_agent_connections = {}
# key: agent_id, value: (reader, writer)

@tweb.route('/')
def index():
    return tweb.Redirect('/index.html')

@tweb.route('/api/agents')
def online_agents():
    # 获取所有在线的agents，agent对象包括agent_id, version, start_time, last_active_time
    agents = list(g_agent_info.values())
    return {'status': 'success', 'data': agents}

# HTTP管理端口 - 配置处理接口
@tweb.route('/config', method='POST')
async def handle_config(request: Request, json_body: JSONBody):
    agent_id = json_body.get('agent_id')
    if not agent_id:
        return Response(status_code=400, body={'error': 'agent_id is required'})
    
    agent_configs[agent_id] = json_body
    return Response(body={'status': 'success', 'message': 'Config saved'})


async def agent_message_dispatcher(reader, writer, agent_id):
    """处理来自agent的消息"""
    while True:
        try:
            type_, value = await decode_tlv(reader)
            if type_ == MSG_TYPE_HEARTBEAT:
                # 接收agent的心跳消息
                stats = json.loads(value.decode())
                log.info(f"[dispatcher][{agent_id}] recv stats: {stats}")
                g_agent_info[agent_id]['total_session'] = stats['total_session']
                continue
            if len(value) < 8:
                log.info(f"[dispatcher][{agent_id}] invalid msg from {agent_id}: missing session_id.")
                continue
            session_id = value[:8].decode()
            # 查找会话
            session = g_sessions.get(session_id)
            if not session:
                log.info(f"[dispatcher][{agent_id}] no session: {session_id}")
                # todo: 通知agent close session
                continue
            if type_ == MSG_TYPE_SESSION_CONNECTED:
                # 标记会话连接成功
                session.target_connected = True
                log.info(f'[dispatcher][{agent_id}][{session_id}] target side connected.')
                # 启动数据转发任务
                session.start_data_forwarder()
            elif type_ == MSG_TYPE_SESSION_DATA:
                if len(value) < 9:
                    log.info(f"[dispatcher][{agent_id}][{session_id}] invalid msg: missing payload.")
                    continue
                payload = value[8:]
                # 放入会话数据队列,session消息循环会自动读取
                await session.tunnel_queue.put(payload)
            elif type_ == MSG_TYPE_SESSION_CLOSE:
                log.info(f"[dispatcher][{agent_id}][{session_id}] recv target close msg.")
                if session.client_writer:
                    session.tunnel_queue.put_nowait(b'')  # empty msg means closed msg.
                    del g_sessions[session_id]  # 清除session数据
            else:
                log.info(f"[dispatcher][{agent_id}] recv unknown msg type: {type_}")
        except Exception as e:
            log.info(f"[dispatcher][{agent_id}] error: {e}")
            break
    log.info(f"[dispatcher][{agent_id}] disconnected. exit loop")

async def start_forward_server(listen_addr, agent_id, target_addr):
    async def handle_client(client_reader, client_writer):
        # 连接到target
        sess = ServerSession(listen_addr, client_reader, client_writer)
        cli_addr = client_writer.get_extra_info("peername")
        sess.client_conn = f'{cli_addr[0]}:{cli_addr[1]} -> {listen_addr}'
        log.info(f'[{sess.session_id}] new conn: {sess.client_conn} -> {agent_id}:{target_addr}')
        sess.type = 'local' if agent_id == 'local' else agent_id
        # 1. local forward scene
        if agent_id == 'local':
            try:
                target_host, target_port = target_addr.split(':')
                target_reader, target_writer = await asyncio.open_connection(target_host, int(target_port))
                sess.target_reader = target_reader
                sess.target_writer = target_writer
                # socket localip:port -> targetip:port
                sess.target_conn = f'{target_writer.get_extra_info("localaddr")} -> {target_addr}'
            except Exception as e:
                log.info(f"无法连接到目标 {target_addr}: {e}")
                client_writer.close()
                # await client_writer.wait_closed()
                return
            # 启动local双向转发
            g_sessions[sess.session_id] = sess
            return sess.start_data_forwarder()
        # 2. tunnel转发场景：获取agent连接
        agent_conn = g_agent_connections.get(agent_id)
        if not agent_conn:
            log.info(f'Agent {agent_id} 未连接')
            client_writer.close()
            # await client_writer.wait_closed()
            return
        _, agent_writer = agent_conn
        sess.agent_writer = agent_writer
        # 通过隧道发起一个连接到target的请求，消息中携带target信息
        # 发送TLV格式session_new消息
        session_new_data = sess.session_id.encode() + json.dumps({'target': target_addr}).encode()
        agent_writer.write(encode_tlv(MSG_TYPE_SESSION_NEW, session_new_data))
        await agent_writer.drain()
        # 等待agent_message_dispatcher处理会话连接确认
        g_sessions[sess.session_id] = sess
        sess.target_conn = f'tunnel({agent_id}) -> {target_addr}'
        return
    try:
        host, port = listen_addr.split(':')
        server = await asyncio.start_server(handle_client, host, int(port))
        g_forward_rules[listen_addr] = {
            'listen': listen_addr,
            'agent_id': agent_id,
            'target': target_addr,
            'server': server
        }
        log.info(f'rule listener started: {listen_addr} -> {agent_id}:{target_addr}')
        return True
    except Exception as e:
        log.info(f"启动转发服务器失败 {listen_addr}: {e}")
        return False

@tweb.route('/api/rules', method='POST')
async def rule_add(request: Request, json_body: JSONBody):
    listen = json_body.get('listen')
    agent_id = json_body.get('agent_id')
    target = json_body.get('target')
    log.info(f'create_forward_rule: {listen}, {agent_id}, {target}')
    if not all([listen, agent_id, target]):
        return Response(status_code=400, body={'error': '缺少必填字段: listen, agent_id, target'})

    if listen in g_forward_rules:
        return Response(status_code=400, body={'error': f'监听地址 {listen} 已存在'})

    # if agent_id not in g_agent_connections:
    #     return Response(status_code=404, body={'error': f'Agent {agent_id} 未连接'})

    if await start_forward_server(listen, agent_id, target):
        save_forward_rules()  # 保存规则
        return Response(status_code=201, body={
            'status': 'success',
            'rule': {'listen': listen, 'agent_id': agent_id, 'target': target}
        })
    else:
        return Response(status_code=500, body={'error': f'无法监听地址 {listen}'})

@tweb.route('/api/rules', method='GET')
def rule_list():
    rules = []
    for rule in g_forward_rules.values():
        status = 'online'
        if rule['agent_id'] != 'local' and rule['agent_id'] not in g_agent_info:
            status = 'offline'
        rules.append({
            'listen': rule['listen'],
            'agent_id': rule['agent_id'],
            'target': rule['target'],
            'status': status
        })
    return {'status': 'success', 'data': rules}

@tweb.route('/api/rules', method='DELETE')
def rule_delete(listen=tweb.Parameter()):
    if listen in g_forward_rules:
        rule = g_forward_rules[listen]
        if rule['server']:
            log.info(f'closing listener {listen}')
            rule['server'].close()
            del g_forward_rules[listen]
        for sess in g_sessions.values():
            if sess.listen == listen:
                sess.tunnel_queue.put_nowait(b'')  # empty msg means closed msg.
        save_forward_rules()  # 保存规则
        return {'status': 'success', 'rule': {'listen': listen, 'agent_id': rule['agent_id'], 'target': rule['target']}}
    else:
        return Response(status_code=404, body={'error': f'监听地址 {listen} 不存在'})
 
@tweb.route('/api/sessions', method='GET')
def session_list():
    sessions = []
    for sess in g_sessions.values():
        sessions.append({
            'sid': sess.session_id,
            'duration': int(time.time() - sess.created_at),
            'bytes_tx': sess.bytes_tx,
            'bytes_rx': sess.bytes_rx,
            'client_conn': sess.client_conn,
            'listen': sess.listen,
            'target_conn': sess.target_conn,
            'type': sess.type,
        })
    return {'status': 'success', 'data': sessions}


@tweb.route('/config/<agent_id>', method='GET')
async def get_config(agent_id: str):
    config = agent_configs.get(agent_id)
    if not config:
        return Response(status_code=404, body={'error': 'Agent not found'})
    return Response(body=config)

@tweb.route('/api/login')
def login(req=tweb.Request()):
    log.info(f'new login request: {req.json}')
    if req.json.get('username') == 'admin' and req.json.get('password') == 'admin':
        token = str(time.time())
        global_tokens[token] = dict(username=req.json.get('username'),
                                    start_time=datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
        return {"status": 'success', "token": token}
    return {"status": 'error', "result": "invalid username or password"}

@tweb.route('/api/userinfo')
def userinfo(req=tweb.Request()):
    token = req.headers.get('Token')
    return {"status": 'success', "data": global_tokens.get(token)}

# 连接端口处理 - 与agent建立连接
async def handle_agent_connection(reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
    agent_id = None
    agent_addr = ':'.join(map(str, writer.get_extra_info('peername')))
    log.info(f'[register] Agent connected from {agent_addr}')
    try:
        # 读取agent注册信息
        try:
            # 读取TLV格式数据
            type_, value = await decode_tlv(reader)
            if type_ != MSG_TYPE_REGISTER:
                raise TLVError(f"Expected register message (type {MSG_TYPE_REGISTER}), got {type_}")
            agent_info = json.loads(value.decode())
            log.info(f'recv register request {agent_info}')
        except TLVError as e:
            error_msg = json.dumps({'error': str(e)}).encode()
            writer.write(encode_tlv(MSG_TYPE_REGISTER_RESP, error_msg))
            await writer.drain()
            return
        agent_id = agent_info.get('agent_id')
        
        if not agent_id:
            error_data = json.dumps({'error': 'agent_id is required'}).encode()
            writer.write(encode_tlv(MSG_TYPE_REGISTER_RESP, error_data))
            await writer.drain()
            return
        
        # 存储agent连接
        g_agent_connections[agent_id] = (reader, writer)
        g_agent_info[agent_id] = {
            'agent_id': agent_id,
            'version': agent_info.get('version'),
            'addr': agent_addr,
            'start_time': datetime.now().strftime("%Y-%m-%d %H:%M:%S")}
        # 返回当前配置
        config = agent_configs.get(agent_id, {})
        # 发送TLV格式响应
        response_data = json.dumps({'status': 'connected', 'config': config}).encode()
        writer.write(encode_tlv(MSG_TYPE_REGISTER_RESP, response_data))
        await writer.drain()
        
        # 启动消息处理任务
        asyncio.create_task(agent_message_dispatcher(reader, writer, agent_id))
        
        # 保持连接，发送保活消息
        while True:
            try:
                # 发送保活消息
                heartbeat_data = json.dumps({'status': 'alive', 'timestamp': time.time()}).encode()
                writer.write(encode_tlv(MSG_TYPE_HEARTBEAT, heartbeat_data))
                await writer.drain()
                await asyncio.sleep(10)
            except Exception as e:
                log.info(f"agent {agent_id} heartbeat failed: {e}")
                break
    except Exception as e:
        log.info(f'Error handling agent: {e}')
    finally:
        log.info(f'Agent {agent_addr} disconnected.')
        log.info(f'agent:{agent_id} closing related sessions')
        for sess in g_sessions.values():
            if sess.type == agent_id:
                log.info(f'[{sess.session_id}] closing client side')
                sess.client_writer.close()
        log.info(f'agent:{agent_id} closing connections')
        if agent_id and agent_id in g_agent_connections:
            del g_agent_connections[agent_id]
            del g_agent_info[agent_id]
            writer.close()
            # await writer.wait_closed()

async def main():
    # 1.启动agent连接服务
    tunnel_server = await asyncio.start_server(handle_agent_connection, '0.0.0.0', 8889)
    addr = tunnel_server.sockets[0].getsockname()
    log.info(f'tunnel server listening on {addr}')
    # 2. 加载转发规则
    await load_forward_rules()
    await asyncio.gather(
        tweb.start_async(host='0.0.0.0', port=8888, resources={"/**": './'}, prefer_coroutine=True),
        tunnel_server.serve_forever())

if __name__ == '__main__':
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        log.info('Server shutdown')